# -*- coding: utf-8 -*-
import json
import asyncio
from typing import List, Tuple, Dict

from loguru import logger
import pandas as pd
from docx import Document

from service.llm_service.deepseek_chat import DeepSeekChat


class ClauseLoader:
    def __init__(self):
        self.model = DeepSeekChat()

    def parse_file_2_paragraph(self, file_path) -> List[str]:
        # 将docx文件分段返回列表
        res_li = []
        doc = Document(file_path)
        for paragraph in doc.paragraphs:
            if paragraph.text == '':
                continue
            res_li.append(paragraph.text.strip())
        return res_li

    def paragraph_splitter(self, paragraphs: List[str]) -> List[str]:
        # 将段落分块
        # 2k字左右一个块
        symbols = ['，', '。', '：', '；', '！', '？', "'"]
        title_li = []
        for p in paragraphs:
            if len(p) != 0:
                if len(p) < 18 and p[-1] not in symbols:
                    title_li.append(p)
        final_splitted_paras = []
        content = paragraphs[0] + ' '  # 标题后面加一个空格，断句，暗示llm这是一个标题，测试后发现划分效果确实更好了
        for i in range(len(paragraphs)):
            if paragraphs[i] not in title_li:  # 正文直接拼接
                add = paragraphs[i] + ' '
                if len(content) > 5000:
                    logger.error(f"content length > 5000")
                    final_splitted_paras.append(content)
                    content = paragraphs[i] + ' '
                else:
                    content += add
            else:  # 遇到标题判断是否够字数了
                if len(content) > 2000:  # 够字数就更新tmp字符串，并且把内容放到最终输出的列表中
                    final_splitted_paras.append(content)
                    content = paragraphs[i] + ' '
                else:
                    add = paragraphs[i] + ' '
                    content += add
        final_splitted_paras.append(content)  # 最后一个chunk在上面的逻辑中不会被放入列表，要额外放一下
        return final_splitted_paras

    async def async_new_clause_parse_llm(self, number: int, prompt: str) -> Tuple[int, List[Dict]]:
        response = await self.model.async_stream_chat(prompt)
        res = response.strip().split('```json')[1].split('```')[0]
        chunks = json.loads(res)['chunks']
        return number + 1, chunks

    async def async_new_clause_chunk_gen(self, paragraph: List[str]):
        prompt_template = '''帮我将如下法规条款内容进行划分成一个个chunk，要求：
1. 每个chunk的字数不超过300字；
2. chunk划分的依据是语义的独立完整，即，chunk中的内容在陈述一件事，并且分点描述的内容要全部包含进去；
3. 即将处理的条款可能存在多级划分，在遵从前两点要求的情况下，对于整体为3级、4级结构的文档，将最低两级的内容进行合并；
4. 最终结果返回json。 json格式要求为：
```json
{
  "chunks": [
    {
      "id": "1",
      "content": "3 术语解释\n1．危险化学品企业（单位），是指危险化学品生产、经营（带有储存设施）企业及取得危险化学品安全使用许可证的企业或单位。"
    },
    {
      "id": "2",
      "content": "2．从业人员，是指从事生产经营活动的所有人员，包括正式员工和劳动派遣人员。"
    },
    ...
    {
       "id": "10",
       "content": "4.2 各部门负责人\n4.2.1 组织制订、审批、实施区域内 HSE 培训计划；\n4.2.2 组织制定并审批区域内员工的培训、再培训课程；\n4.2.3 定期组织对员工能力进行评估；\n4.2.4 确保区域内员工接受国家法规要求和公司要求的培训。"
    }
  ]
}
```
待划分的法规条款为：'''

        tasks = [self.async_new_clause_parse_llm(_no, prompt_template + p) for _no, p in enumerate(paragraph)]
        logger.info(f"task start  len:{len(tasks)}")
        async_results: List[Tuple[int, List[Dict]]] = await asyncio.gather(*tasks)
        logger.info(f"task end:{async_results}")
        sorted_results = sorted(async_results, key=lambda x: x[0])
        c = 0
        final_res = []
        for element in sorted_results:
            for chunk in element[1]:
                c += 1
                final_res.append({'line_number': c, 'content': chunk['content']})
        return final_res

    def new_clause_chunk_gen(self, paragraph: List[str]):
        prompt_template = '''帮我将如下法规条款内容进行划分成一个个chunk，要求：
1. 每个chunk的字数不超过300字；
2. chunk划分的依据是语义的独立完整，即，chunk中的内容在陈述一件事，并且分点描述的内容要全部包含进去；
3. 即将处理的条款可能存在多级划分，在遵从前两点要求的情况下，对于整体为3级、4级结构的文档，将最低两级的内容进行合并；
4. 最终结果返回json。 json格式要求为：
```json
{
  "chunks": [
    {
      "id": "1",
      "content": "3 术语解释\n1．危险化学品企业（单位），是指危险化学品生产、经营（带有储存设施）企业及取得危险化学品安全使用许可证的企业或单位。"
    },
    {
      "id": "2",
      "content": "2．从业人员，是指从事生产经营活动的所有人员，包括正式员工和劳动派遣人员。"
    },
    ...
    {
       "id": "10",
       "content": "4.2 各部门负责人\n4.2.1 组织制订、审批、实施区域内 HSE 培训计划；\n4.2.2 组织制定并审批区域内员工的培训、再培训课程；\n4.2.3 定期组织对员工能力进行评估；\n4.2.4 确保区域内员工接受国家法规要求和公司要求的培训。"
    }
  ]
}
```
待划分的法规条款为：'''
        c = 0
        final_res = []
        for p in paragraph:
            res = ''  # 一个生成的json字符串
            for i in self.model.stream_chat(prompt_template + p):
                res += i
            res = res.strip().split('```json')[1].split('```')[0]
            chunks = json.loads(res)['chunks']
            for chunk in chunks:
                c += 1
                final_res.append({'line_number': c, 'content': chunk['content']})
        return final_res

    async def load_docx(self, file_path):
        parsed_docx = self.parse_file_2_paragraph(file_path)
        splittered_docx = self.paragraph_splitter(parsed_docx)
        # 划分效果测试
        logger.info(f"split {file_path}, len:{len(splittered_docx)}")
        chunks = await self.async_new_clause_chunk_gen(splittered_docx)
        return chunks

    def read_by_sheet_name(self, file_path: str, clause_name: str = None):
        meta_data = pd.read_excel(file_path)
        standard_clauses = []
        clause_elements = []  # 制度要素列表
        last_element = None
        for i in range(len(meta_data)):
            current_element = meta_data.iloc[i, 0]
            if not pd.isna(current_element) and current_element not in clause_elements:
                clause_elements.append(current_element)
                last_element = current_element
            clause_dict = {
                "line_number": i + 1,
                "clause_name": clause_name,
                'institution': last_element if pd.isna(current_element) else current_element,
                'clause_number': None if pd.isna(meta_data.iloc[i, 1]) else meta_data.iloc[i, 1],
                'original_text': None if pd.isna(meta_data.iloc[i, 2]) else meta_data.iloc[i, 2],
                'artificial_text': None if pd.isna(meta_data.iloc[i, 3]) else meta_data.iloc[i, 3],
                'area': None if pd.isna(meta_data.iloc[i, 4]) else meta_data.iloc[i, 4],
                'cargo': None if pd.isna(meta_data.iloc[i, 5]) else meta_data.iloc[i, 5],
            }
            standard_clauses.append(clause_dict)
        return standard_clauses

    def standard_clause_parse(self, clause_name: str = None,
                              file_path: str = '../01 安全生产教育培训制度等——合规性审核数据表.xlsx',
                              ):
        return self.read_by_sheet_name(file_path, clause_name)

